package org.bilinc.gz001;

import javax.swing.*;
import javax.swing.plaf.basic.BasicTableHeaderUI;
import javax.swing.table.TableCellRenderer;
import javax.swing.table.TableColumn;
import javax.swing.table.TableColumnModel;
import java.awt.*;
import java.util.Enumeration;
import java.util.Hashtable;

public class GroupableTableHeaderUI extends BasicTableHeaderUI {

    public void paint(Graphics g, JComponent c) {
        final Rectangle clipBounds = g.getClipBounds();

        if (header.getColumnModel() == null) {
            return;
        }

        ((GroupableTableHeader) header).setColumnMargin();
        int column = 0;
        final Dimension size = header.getSize();
        final Rectangle cellRect = new Rectangle(0, 0, size.width, size.height);
        final Hashtable<ColumnGroup, Rectangle> h = new Hashtable<ColumnGroup, Rectangle>();

        final Enumeration<TableColumn> enumeration = header.getColumnModel().getColumns();
        while (enumeration.hasMoreElements()) {
            cellRect.height = size.height;
            cellRect.y = 0;
            final TableColumn aColumn = enumeration.nextElement();
            final Enumeration cGroups = ((GroupableTableHeader) header).getColumnGroups(aColumn);
            if (cGroups != null) {
                int groupHeight = 0;
                while (cGroups.hasMoreElements()) {
                    final ColumnGroup cGroup = (ColumnGroup) cGroups.nextElement();
                    Rectangle groupRect = h.get(cGroup);
                    if (groupRect == null) {
                        groupRect = new Rectangle(cellRect);
                        Dimension d = cGroup.getSize(header.getTable());
                        groupRect.width = d.width;
                        groupRect.height = d.height;
                        h.put(cGroup, groupRect);
                    }

                    paintCell(g, groupRect, cGroup);
                    groupHeight += groupRect.height;
                    cellRect.height = size.height - groupHeight;
                    cellRect.y = groupHeight;
                }
            }

            cellRect.width = aColumn.getWidth();
            if (cellRect.intersects(clipBounds)) {
                paintCell(g, cellRect, column);
            }

            cellRect.x += cellRect.width;
            column++;
        }
    }


    private void paintCell(final Graphics g, final Rectangle cellRect, final int columnIndex) {
        final TableColumn aColumn = header.getColumnModel().getColumn(columnIndex);
        final TableCellRenderer renderer = header.getDefaultRenderer();

        final Component component = renderer.getTableCellRendererComponent(header.getTable(), aColumn
                .getHeaderValue(), false, false, -1, columnIndex);

        rendererPane.add(component);

        rendererPane.paintComponent(g, component, header, cellRect.x, cellRect.y, cellRect.width,
                cellRect.height, true);
    }


    private void paintCell(final Graphics g, final Rectangle cellRect, final ColumnGroup cGroup) {
        final  TableCellRenderer renderer = cGroup.getHeaderRenderer();
        final Component component = renderer.getTableCellRendererComponent(header.getTable(), cGroup
                .getHeaderValue(), false, false, -1, -1);
        component.setBackground(Color.PINK);
        rendererPane.add(component);

        rendererPane.paintComponent(g, component, header, cellRect.x, cellRect.y, cellRect.width,
                cellRect.height, true);
    }


    private int getHeaderHeight() {
        int height = 0;
        final TableColumnModel columnModel = header.getColumnModel();

        for (int column = 0; column < columnModel.getColumnCount(); column++) {
            final TableColumn aColumn = columnModel.getColumn(column);
            final TableCellRenderer renderer = header.getDefaultRenderer();

            final Component comp = renderer.getTableCellRendererComponent(header.getTable(), aColumn
                    .getHeaderValue(), false, false, -1, column);
            int cHeight = comp.getPreferredSize().height;
            final Enumeration enumeration = ((GroupableTableHeader) header).getColumnGroups(aColumn);
            if (enumeration != null) {
                while (enumeration.hasMoreElements()) {
                    ColumnGroup cGroup = (ColumnGroup) enumeration.nextElement();
                    cHeight += cGroup.getSize(header.getTable()).height;
                }
            }
            height = Math.max(height, cHeight);
        }
        return height;
    }


    private Dimension createHeaderSize(long width) {
        final TableColumnModel columnModel = header.getColumnModel();
        width += columnModel.getColumnMargin() * columnModel.getColumnCount();
        if (width > Integer.MAX_VALUE) {
            width = Integer.MAX_VALUE;
        }
        return new Dimension((int) width, getHeaderHeight());
    }


    public Dimension getPreferredSize(final JComponent c) {
        long width = 0;
        final Enumeration<TableColumn> enumeration = header.getColumnModel().getColumns();
        while (enumeration.hasMoreElements()) {
            final TableColumn aColumn = enumeration.nextElement();
            width = width + aColumn.getPreferredWidth();
        }
        return createHeaderSize(width);
    }
}